from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, IntegerType, StringType

if __name__ == '__main__':
    # 构建SparkSession对象
    spark = SparkSession.builder. \
        appName("local[*]"). \
        config("spark.sql.shuffle.partitions", "4"). \
        getOrCreate()
    # appName 设置程序名称
    # config 设置常用属性。可以通过此来设置配置
    # 最后通过getOrCreate 创建 SparkSession对象

    # 从SparkSession中获取SparkContext
    sc = spark.sparkContext

    # 1.获取数据集
    schema = StructType() \
        .add("user_id", StringType(), nullable=True) \
        .add("movie_id", IntegerType(), nullable=True) \
        .add("rank", IntegerType(), nullable=True) \
        .add("ts", StringType(), nullable=True)

    df = spark.read.format("csv") \
        .option("sep", "\t") \
        .option("header", "false") \
        .option("encoding", "utf-8") \
        .schema(schema) \
        .load("hdfs://node1:8020/input/u.data")

    # TODO 1:写出TEXT文件
    # 写出TEXT，需要转换成一列写出
    df.select(F.concat_ws("-----", "user_id", "movie_id", "rank", "ts")) \
        .write \
        .format("text") \
        .mode("overwrite") \
        .save("../../data/output/text")

    # TODO 2:写出CSV文件
    df.write \
        .format("csv") \
        .mode("overwrite") \
        .option("sep", ",") \
        .option("header", "true") \
        .option("encoding", "utf-8") \
        .save("../../data/output/csv")

    # TODO 3:写出JSON文件
    df.write \
        .format("json") \
        .mode("overwrite") \
        .save("../../data/output/json")

    # TODO 4:写出parquet文件
    df.write \
        .format("parquet") \
        .mode("overwrite") \
        .save("../../data/output/parquet")
